%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Explanation of what the file does
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% This function solves for mobility matrices across periods and a labor 
% supply matrix across periods given a guess for the path of utilities
% as well as initial mobility matrices and an initial labor supply
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Inputs needed
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% 1. Guess of path of utility dots (Ugu), size I*S*T
% 2. Initial mobility matrix (MuInit), size (I*S)*(I*S)
% 3. Initial labor supply matrix (EllInit), size I*S
% 4. Discount factor (ba), scalar
% 5. Migration elasticity (ka), scalar
% 6. Sectoral elasticity (nu), scalar
% 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Outputs produced
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% 1. Mobility matrix (MuEra), size (I*S)*(I*S)*T
% 2. Labor supply matrix (EllEra), size I*S*T

function [MuEra,EllEra] = NEBA_Block1(Ugu,MuInit,EllInit,ba,ka,nu)

% Define some important variables

I                 = size(EllInit,1);
S                 = size(EllInit,2);
T                 = size(Ugu,3);
MuEra             = NaN(I*S,I*S,T);
EllEra            = NaN(I,S,T);
EllEra(:,:,1)     = EllInit;

% Notice that MuEra, EllEra and Uguess all are of size T in the time
% dimension, but they encompass different periods. Mu and Ell go from
% period zero to period T-1, while Uguess goes from period one to period T

% Compute Mu for period zero

UDI               = Ugu(:,:,1)';
MuHashInit        = MuInit * kron(eye(I),ones(S,1));
MuCondInit        = MuInit ./ kron(MuHashInit,ones(1,S));
MuCondInit(isnan(MuCondInit)) = 0;
MuCondNum         = (repmat(UDI(:)',I*S,1)).^(ba/nu) .* MuCondInit;
MuCondDen         = kron(MuCondNum*kron(eye(I),ones(S,1)),ones(1,S));
MuCond            = MuCondNum./MuCondDen;
MuCond(isnan(MuCond)) = 0;
MuHashNum         = MuHashInit .* (MuCondNum*kron(eye(I),ones(S,1)))...
                    .^(nu/ka);
MuHashDen         = kron(sum(MuHashNum,2),ones(1,I));
MuHash            = MuHashNum./MuHashDen;
MuEra(:,:,1)      = MuCond .* kron(MuHash,ones(1,S));

% Run the forward loop where the mus of a period are found using the mus of
% the previous period and the u dots of the following period and also
% obtain labor supplies from previous period labor supplies and mus

for t=2:T

    UDI           = Ugu(:,:,t)';
    
    MuHashPre     = MuEra(:,:,t-1) * kron(eye(I),ones(S,1));
    MuCondPre     = MuEra(:,:,t-1) ./ kron(MuHashPre,ones(1,S));
    MuCondPre(isnan(MuCondPre)) = 0;
    MuCondNum     = (repmat(UDI(:)',I*S,1)).^(ba/nu) .* MuCondPre;
    MuCondDen     = kron(MuCondNum*kron(eye(I),ones(S,1)),ones(1,S));
    MuCond        = MuCondNum./MuCondDen;
    MuCond(isnan(MuCond)) = 0;
    MuHashNum     = MuHashPre .* (MuCondNum*kron(eye(I),ones(S,1)))...
                        .^(nu/ka);
    MuHashDen     = kron(sum(MuHashNum,2),ones(1,I));
    MuHash        = MuHashNum./MuHashDen;
    MuEra(:,:,t)  = MuCond .* kron(MuHash,ones(1,S));
    
    EllInter      = EllEra(:,:,t-1)';
    EllNewV       = MuEra(:,:,t-1)' * EllInter(:);
    EllEra(:,:,t) = reshape(EllNewV,S,I)';

end

end
